#include "meshTypeFunctions.hpp"
#include "structMeshPartitioning.hpp"
#include "mpiH_Comm.hpp"
#include "mpiH_Func.hpp"
#include <catch2/catch_test_macros.hpp>
#include <array>
#include <cmath>

using namespace hms;

bool operator==( const OverlapSegment& a, const OverlapSegment b );
std::array<Array2Xi, 6> sizesAndOffsets();

TEST_CASE( MESH_NAME ": Constructing StructMeshPart in different ways" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	if ( comm.size() != 12 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: Most tests relating to the StructMeshPart "
				"class will not be executed if the distributed computing tests "
				"are executed with a number of processes other than 12 "
				"(i.e. the -np argument for mpirun).\n";
	} else {
		Array2i globalMeshSize {16,12};

		auto test = [&]( const StructMeshPart<MESH_TYPE>& part ){
			REQUIRE( comm.size() == part.comm().size() );
			REQUIRE( comm.rank() == part.comm().rank() );

			REQUIRE( (part.meshSizeGlobal() == globalMeshSize).all() );
			Array2Xi patchLimits (2, 5);
			patchLimits <<
				0,  6, 208, 224, 192,
				5, 15, 395, 411, 207;
			REQUIRE( (part.boundaryPatchLimits() == patchLimits).all() );

			if ( comm.size() == 12 ){
				REQUIRE( (part.meshSize() == Array2i{4,4}).all() );
				REQUIRE( (part.nParts() == Array2i{4,3} ).all() );
				REQUIRE( part.band() == comm.rank() / 4 );
				REQUIRE( part.indexInBand() == comm.rank() % 4 );
				REQUIRE( (part.offset() == Array2i{
					part.indexInBand() * 4, part.band() * 4
				}).all() );

				if ( comm.rank() < 8 )
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{ comm.rank() + 4, 4, 0, 0 } );
				if ( comm.rank() >= 4 )
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{ comm.rank() - 4, 4, 0, 0 } );
				
				if ( comm.rank() < 4 )
					REQUIRE( part.atBoundary(BoundaryLocation::bottom) );
				if ( comm.rank() >= 8 )
					REQUIRE( part.atBoundary(BoundaryLocation::top) );
				if ( comm.rank() % 4 == 0 )
					REQUIRE( part.atBoundary(BoundaryLocation::left) );
				if ( comm.rank() % 4 == 3 )
					REQUIRE( part.atBoundary(BoundaryLocation::right) );
			}
		};

		auto addGlobalPatches = [&]( MESH_TYPE& mesh ){
			addStandardBoundaryPatches(mesh);

			/* change up boundary patches a bit */
			mesh.eraseBoundaryPatch("bottom");

			mesh.addBoundaryPatch(
				"bottomLeft", ArrayXi::LinSpaced( 6, 0, 5 )
			);
			mesh.addBoundaryPatch(
				"bottomRight", ArrayXi::LinSpaced( 10, 6, 15 )
			);
		};

		auto addLocalPatches = [&]( MESH_TYPE& mesh ){
			addStandardBoundaryPatches(mesh);

			if ( comm.rank() < 8 )
				mesh.boundaryPatch("top") = BoundaryPatch{"top"};
			if ( comm.rank() % 4 != 0 )
				mesh.boundaryPatch("left") = BoundaryPatch{"left"};
			if ( comm.rank() % 4 != 3 )
				mesh.boundaryPatch("right") = BoundaryPatch{"right"};

			/* change up boundary patches a bit */
			mesh.eraseBoundaryPatch("bottom");
			mesh.addBoundaryPatch( "bottomLeft", ArrayXi{} );
			mesh.addBoundaryPatch( "bottomRight", ArrayXi{} );

			if ( comm.rank() == 0 )
				mesh.boundaryPatch("bottomLeft").setEdgesAndCells(
					mesh, ArrayXi::LinSpaced(4, 0, 3) );
			if ( comm.rank() == 1 ){
				mesh.boundaryPatch("bottomLeft").setEdgesAndCells(
					mesh, ArrayXi::LinSpaced(2, 0, 1) );
				mesh.boundaryPatch("bottomRight").setEdgesAndCells(
					mesh, ArrayXi::LinSpaced(2, 2, 3) );
			}
			if ( comm.rank() == 2 || comm.rank() == 3 )
				mesh.boundaryPatch("bottomRight").setEdgesAndCells(
					mesh, ArrayXi::LinSpaced(4, 0, 3) );
		};

		auto localOffset = [&]() -> Array2i {
			return { (comm.rank() % 4) * 4, (comm.rank() / 4) * 4 };
		};

		SECTION( "Default initialisation, then assigning members" ){

			StructMeshPart<MESH_TYPE> part;
			part.comm( comm );

			SECTION( "Assigning global mesh on master, then distributing" ){
				auto assignGlobal = [&](){

					MESH_TYPE mesh;

					if ( comm.isMaster() ){
						mesh = uniMeshAs<MESH_TYPE>( globalMeshSize, (Array2s{1,1}),
							(Array2s{0,0}) );
						addGlobalPatches(mesh);
					}
					part.assignAndDistribute( std::move(mesh) );

					test(part);
				};
				assignGlobal();
			}
			SECTION( "Assigning local mesh" ){
				auto assignLocal = [&](){
					MESH_TYPE mesh { uniMeshAs<MESH_TYPE>( (Array2i{4,4}), (Array2s{1,1}),
						(Array2s{0,0}) ) };
					addLocalPatches(mesh);

					part.assignLocal( std::move(mesh), localOffset() );

					test(part);
				};
				assignLocal();
			}
		}
		SECTION( "Assigning mesh during construction" ){
			SECTION( "Assigning global mesh on master" ){
				auto assignGlobal = [&](){
					MESH_TYPE mesh;

					if ( comm.isMaster() ){
						mesh = uniMeshAs<MESH_TYPE>( globalMeshSize, (Array2s{1,1}),
							(Array2s{0,0}) );
						addGlobalPatches(mesh);
					}
					StructMeshPart part ( std::move(mesh), comm );

					test(part);
				};
				assignGlobal(1);
			}
			SECTION( "Assigning local meshes" ){
				auto assignLocal = [&](){
					MESH_TYPE mesh { uniMeshAs<MESH_TYPE>( (Array2i{4,4}), (Array2s{1,1}),
						(Array2s{0,0}) ) };
					addLocalPatches(mesh);

					StructMeshPart part { std::move(mesh), localOffset(), comm };
					
					test(part);
				};
				assignLocal();
			}
		}
	}
}


/* This test is very easy to cause deadlocks. If there is a deadlock and the
 * program hangs, just try commenting out this test and see whether it's 
 * working again.
 */
TEST_CASE( MESH_NAME ": Redistributing mesh" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	auto test = [&]( int nGhostLayers ){
		if ( comm.size() != 12 ) return;
		
		Array2i globalMeshSize {17,13};

		auto [meshSizes0, meshSizes1, meshSizes2, offsets0, offsets1, offsets2]
			= sizesAndOffsets();

		MESH_TYPE mesh;

		if ( comm.isMaster() ){
			mesh = uniMeshAs<MESH_TYPE>( globalMeshSize, (Array2s{1,1}),
				(Array2s{0,0}), nGhostLayers );
			addStandardBoundaryPatches(mesh);
		}
			
		StructMeshPart part { std::move(mesh), comm };

		REQUIRE( (part.meshSize() == meshSizes0.col( comm.rank() ) ).all() );
		REQUIRE( (part.offset() == offsets0.col( comm.rank() ) ).all() );

		auto testStage1 = [&, &part = std::as_const(part)](){
			REQUIRE( (part.meshSize() == meshSizes1.col( comm.rank() ) ).all() );
			REQUIRE( (part.offset() == offsets1.col( comm.rank() ) ).all() );
			REQUIRE( (part.oldSizes() == meshSizes0).all() );

			switch( comm.rank() ){
				case 0:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{4, 3, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{5, 2, 0, 3} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(5, 0, 4) ).all() );
					REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
						ArrayXi::LinSpaced(5, 30, 54) ).all() );
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 1:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{5, 2, 2, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{6, 2, 0, 2} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(4, 0, 3) ).all() );
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 2:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{6, 4, 2, 0} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(4, 0, 3) ).all() );
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 3:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{7, 4, 0, 0} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(4, 0, 3) ).all() );
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE( (part.mesh().boundaryPatch("right").edges() ==
						ArrayXi::LinSpaced(5,28,48) ).all() );
					break;
				case 4:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{0, 3, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{8, 3, 0, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
						ArrayXi::LinSpaced(2, 9,13) ).all() );
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 5:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{0, 2, 3, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{1, 2, 0, 2} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{8, 1, 3, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{9, 3, 0, 1} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 6:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{ 1, 2, 2, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{ 2, 4, 0, 2} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{10, 4, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{11, 2, 0, 4} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 7:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{ 3, 4, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{11, 4, 2, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("right").edges() ==
						ArrayXi::LinSpaced(2,16,21) ).all() );
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					break;
				case 8:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{4, 3, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{5, 1, 0, 3} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
						ArrayXi::LinSpaced(6,28,53) ).all() );
					REQUIRE(part.mesh().boundaryPatch("right" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(4,24,27) ).all() );
					break;
				case 9:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{5, 3, 1, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("right" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(3,18,20) ).all() );
					break;
				case 10:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{6, 4, 0, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("right" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(4,24,27) ).all() );
					break;
				case 11:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{6, 2, 4, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{7, 4, 0, 2} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("right").edges() ==
						ArrayXi::LinSpaced(6,48,83) ).all() );
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(6,36,41) ).all() );
					break;
			};
		};

		auto testStage2 = [&, &part = std::as_const(part)](){
			REQUIRE( (part.meshSize() == meshSizes2.col( comm.rank() ) ).all() );
			REQUIRE( (part.offset() == offsets2.col( comm.rank() ) ).all() );
			REQUIRE( (part.oldSizes() == meshSizes1).all() );

			switch( comm.rank() ){
				case 0:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{4, 3, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{5, 2, 0, 3} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(5, 0, 4) ).all() );
					REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
						ArrayXi::LinSpaced(5, 30, 54) ).all() );
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 1:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{5, 2, 2, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{6, 2, 0, 2} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(4, 0, 3) ).all() );
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 2:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{6, 3, 2, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{7, 1, 0, 3} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(4, 0, 3) ).all() );
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 3:
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{7, 4, 1, 0} );

					REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
						ArrayXi::LinSpaced(4, 0, 3) ).all() );
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE( (part.mesh().boundaryPatch("right").edges() ==
						ArrayXi::LinSpaced(5,28,48) ).all() );
					break;
				case 4:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{0, 3, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{8, 3, 0, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
						ArrayXi::LinSpaced(3,12,20) ).all() );
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 5:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{0, 2, 3, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{1, 2, 0, 2} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{8, 1, 3, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{9, 3, 0, 1} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 6:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{ 1, 2, 2, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{ 2, 3, 0, 2} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{10, 4, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[1] ==
						OverlapSegment{11, 1, 0, 4} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					REQUIRE(part.mesh().boundaryPatch("right").empty());
					break;
				case 7:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{ 2, 1, 3, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{ 3, 4, 0, 1} );
					REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==
						OverlapSegment{11, 5, 1, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("right").edges() ==
						ArrayXi::LinSpaced(3,25,37) ).all() );
					REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
					break;
				case 8:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{4, 3, 0, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{5, 1, 0, 3} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
						ArrayXi::LinSpaced(5,24,44) ).all() );
					REQUIRE(part.mesh().boundaryPatch("right" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(4,20,23) ).all() );
					break;
				case 9:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{5, 3, 1, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("right" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(3,15,17) ).all() );
					break;
				case 10:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{6, 4, 0, 0} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE(part.mesh().boundaryPatch("right" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(4,20,23) ).all() );
					break;
				case 11:
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] ==
						OverlapSegment{6, 1, 4, 0} );
					REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] ==
						OverlapSegment{7, 5, 0, 1} );

					REQUIRE(part.mesh().boundaryPatch("bottom" ).empty());
					REQUIRE(part.mesh().boundaryPatch("left" ).empty());
					REQUIRE( (part.mesh().boundaryPatch("right").edges() ==
						ArrayXi::LinSpaced(5,42,70) ).all() );
					REQUIRE( (part.mesh().boundaryPatch("top").edges() ==
						ArrayXi::LinSpaced(6,30,35) ).all() );
					break;
			}
		};

		SECTION( "Pure redistribution - "
			+ std::to_string(nGhostLayers) + " ghost layers"
		){
			part.redistributeMesh( meshSizes1 );
					
			SECTION( "Single redistribution" ){
				testStage1();
			}
			SECTION( "Repeated redistribution" ){
				part.redistributeMesh( meshSizes2 );
				testStage2();
			}
		}
		SECTION( "Rebalancing - "
			+ std::to_string(nGhostLayers) + " ghost layers"
		){
			ArrayXs loadPerNode1, loadPerNode2;
			loadPerNode1.resize( comm.size() );
			loadPerNode2.resizeLike(loadPerNode1);
			loadPerNode1 <<
				0.1, 0.2, 0.4, 0.3, 1.0, 0.8, 0.5, 0.7, 0.2, 0.4, 0.2, 0.1;

			loadPerNode2 <<
				0.288, 0.345, 0.41, 0.38, 1.0, 0.85, 7./12, 0.7, 1./3, 7./15, 73./240, 25./120;
			
			SECTION( "by local load per cell"){
				ArrayXs load ( part.meshSize().prod() );
				load = loadPerNode1( comm.rank() );

				part.rebalanceMeshSizes( PartLoad::perCell, load );
				SECTION( "Single redistribution" ){
					testStage1();
				}
				SECTION( "Repeated redistribution" ){
					load.resize( part.meshSize().prod() );
					load = loadPerNode2( comm.rank() );

					part.rebalanceMeshSizes( PartLoad::perCell, load );
					testStage2();
				}
			}
			SECTION( "by global load per cell" ){
				ArrayXXs load (
					part.meshSizeGlobal().x(), part.meshSizeGlobal().y() );
				for ( int i{0}; i<comm.size(); ++i ){
					load.block(
						offsets0(0,i), offsets0(1,i),
						meshSizes0(0,i), meshSizes0(1,i)
					) = loadPerNode1(i);
				}
				part.rebalanceMeshSizes(
					PartLoad::perCellRecomp, load.reshaped() );
				SECTION( "Single redistribution" ){
					testStage1();
				}
				SECTION( "Repeated redistribution" ){
					for ( int i{0}; i<comm.size(); ++i ){
						load.block(
							offsets1(0,i), offsets1(1,i),
							meshSizes1(0,i), meshSizes1(1,i)
						) = loadPerNode2(i);
					}
					part.rebalanceMeshSizes(
						PartLoad::perCellRecomp, load.reshaped() );
					testStage2();
				}
			}
			SECTION( "by load per node" ){
				part.rebalanceMeshSizes( PartLoad::perNode, loadPerNode1 );
				SECTION( "Single redistribution" ){
					testStage1();
				}
				SECTION( "Repeated redistribution" ){
					part.rebalanceMeshSizes( PartLoad::perNode, loadPerNode2 );
					testStage2();
				}
			}
		}
	};
	test(1);
	test(2);
}

/* fields are not yet implemented fully, so these tests should be revisited */
TEST_CASE( MESH_NAME ": (Re-)Distributing field" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	if ( comm.size() == 12 ){
		Array2i globalMeshSize {17,13};

		auto [meshSizes0, meshSizes1, meshSizes2, offsets0, offsets1, offsets2]
			= sizesAndOffsets();
		
		MESH_TYPE mesh;
		Field<MESH_TYPE> f;

		auto writeRanksToField = [&](
			int nComponents, int ng,
			const Ref<const Array2Xi>& offsets,
			const Ref<const Array2Xi>& meshSizes
		){
			/* write the rank who will get that part of the field into its
			 * values, so that each rank can check the values */
			if ( comm.isMaster() ){
				mesh = uniMeshAs<MESH_TYPE>(
					globalMeshSize, (Array2s{1,1}), (Array2s{0,0}), ng
				);
				f = Field( mesh, nComponents );
				f.values() = 0.99;
				for ( int i{0}; i<comm.size(); ++i )
					f.values().reshaped(
						(mesh.nCellsX() + 2*ng)*nComponents,
						mesh.nCellsY() + 2*ng
					).block(
						(offsets(0,i)+ng)*nComponents, offsets(1,i)+ng,
						meshSizes(0,i)*nComponents, meshSizes(1,i)
					) = i;
			}
		};

		SECTION( "Distributing field" ){
			auto test = [&]( int nComponents, int ng ){
				writeRanksToField( nComponents, ng, offsets0, meshSizes0 );

				StructMeshPart part ( std::move(mesh), comm );

				part.distributeField( f );

				REQUIRE( ( f.domainValues() == comm.rank() ).all() );
			};

			for ( int nVals {1}; nVals<4; ++nVals )
				for ( int ng {1}; ng<3; ++ng )
					test( nVals, ng );
		}
		SECTION( "Redistributing field once" ){
			auto test = [&]( int nComponents, int ng ){
				writeRanksToField( nComponents, ng, offsets1, meshSizes1 );
				
				StructMeshPart part ( std::move(mesh), comm );
				part.distributeField(f);
				part.redistributeMesh( meshSizes1 );
				part.redistributeField(f);

				REQUIRE( ( f.domainValues() == comm.rank() ).all() );
			};
			for ( int nVals {1}; nVals<4; ++nVals )
				for ( int ng {1}; ng<3; ++ng )
					test( nVals, ng );
		}
		SECTION( "Repeatedly redistributing field" ){
			auto test = [&](int nComponents, int ng){
				writeRanksToField( nComponents, ng, offsets2, meshSizes2 );

				StructMeshPart part ( std::move(mesh), comm );
				part.distributeField(f);
				part.redistributeMesh( meshSizes1 );
				part.redistributeField(f);
				part.redistributeMesh( meshSizes2 );
				part.redistributeField(f);
				
				REQUIRE( ( f.domainValues() == comm.rank() ).all() );
			};
			for ( int nVals {1}; nVals<4; ++nVals )
				for ( int ng {1}; ng<3; ++ng )
					test( nVals, ng );
		}
	}
}

TEST_CASE( MESH_NAME ": getNeighbourValues: Retrieving domain values from "
	"neighbours to update ghost values"
){
	MPI_Comm_H comm { MPI_COMM_WORLD };
	Array2i meshSizeGlobal {30,20};

	auto test = [&]( Index nGhostLayers, Index nComponents ){

		MESH_TYPE mesh;

		StructMeshPart part {
			comm.isMaster() ?
				uniMeshAs<MESH_TYPE>( meshSizeGlobal, (Array2s{1,1}), (Array2s{0,0}),
					nGhostLayers ) :
				std::move(mesh),
			comm
		};

		Field f { part.mesh(), nComponents };

		auto testField = [&](){

			f.values() = -1;

			f.domainValues() = comm.rank();
			part.getNeighbourValues( f );

			auto valsReshaped { f.values().reshaped(
				(part.meshSize().x() + 2 * nGhostLayers)*nComponents,
				part.meshSize().y() + 2 * nGhostLayers
			) };

			auto bottomVals { valsReshaped.block(
				nGhostLayers*nComponents, 0,
				part.meshSize().x()*nComponents, nGhostLayers
			) };
			auto topVals { valsReshaped.block(
				nGhostLayers*nComponents, nGhostLayers + part.meshSize().y(),
				part.meshSize().x()*nComponents, nGhostLayers
			) };
			auto leftVals { valsReshaped.block(
				0, nGhostLayers, nGhostLayers*nComponents, part.meshSize().y()
			) };
			auto rightVals { valsReshaped.block(
				(nGhostLayers + part.meshSize().x())*nComponents, nGhostLayers,
				nGhostLayers*nComponents, part.meshSize().y()
			) };

			for ( const auto& nb : std::as_const(part).neighbours(BoundaryLocation::bottom) )
				REQUIRE( bottomVals.middleRows(
					nb.ownBeg*nComponents, nb.size*nComponents
				).isApproxToConstant(nb.rank, epsilon) );
			
			for ( const auto& nb : std::as_const(part).neighbours(BoundaryLocation::top) )
				REQUIRE( ( topVals.middleRows(
					nb.ownBeg*nComponents, nb.size*nComponents
				) == nb.rank ).all() );
			
			if ( !part.atBoundary(BoundaryLocation::left) )
				REQUIRE( (leftVals == comm.rank()-1).all() );
			
			if ( !part.atBoundary(BoundaryLocation::right) )
				REQUIRE( (rightVals == comm.rank()+1).all() );
		};

		testField();
		
		/* create some non-uniform load for rebalancing */
		ArrayXs load ( comm.size() );
		for ( int i{0}; i<comm.size(); ++i ){
			load(i) = comm.size() + i;
		}

		part.rebalanceMeshSizes( PartLoad::perNode, load );

		part.redistributeField(f);

		testField();
	};
	for ( int nVals {1}; nVals<4; ++nVals )
		for ( int ng {1}; ng<3; ++ng )
			test( nVals, ng );
}

TEST_CASE( MESH_NAME ": nRanks invariance, construction", "[any_np]" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };
	Array2i globalMeshSize {240,360};
	int nGhostLayers {2};
	Array2Xi patchLimits (2, 4);
	patchLimits <<
		  0,  86640,  86880 ,86400,
		239, 173159, 173399 ,86639;

	// auto print = [&]( const auto& message ){
	// 	std::stringstream sstream;
	// 	sstream << "Rank " << comm.rank() << ": " << message << "\n";
	// 	for (int i{0};i<comm.size();++i){
	// 		MPI_Barrier(comm);
	// 		if (comm.rank()==i)
	// 			std::cerr << sstream.str();
	// 	}
	// };

	StructMeshPart<MESH_TYPE> part;

	auto test = [&](){
		REQUIRE( comm.size() == part.comm().size() );
		REQUIRE( comm.rank() == part.comm().rank() );

		REQUIRE( (part.meshSizeGlobal() == globalMeshSize).all() );
		REQUIRE( (part.boundaryPatchLimits() == patchLimits).all() );
		REQUIRE( part.meshSize().x() == globalMeshSize.x() / part.nParts().x() +
			(part.indexInBand() == part.nParts().x() - 1 ?
			globalMeshSize.x() % part.nParts().x() : 0)
		);
		REQUIRE( part.meshSize().y() == globalMeshSize.y() / part.nParts().y() +
			(part.band() == part.nParts().y() - 1 ?
			globalMeshSize.y() % part.nParts().y() : 0)
		);
		if ( part.band() == 0 )
			REQUIRE( (part.mesh().boundaryPatch("bottom").edges() ==
				ArrayXi::LinSpaced( part.meshSize().x(), 0, part.meshSize().x()-1)
				).all() );
		else
			REQUIRE( part.mesh().boundaryPatch("bottom").size() == 0 );
		if ( part.indexInBand() == 0 )
			REQUIRE( (part.mesh().boundaryPatch("left").edges() ==
				ArrayXi::LinSpaced(
					part.meshSize().y(),
					part.meshSize().x()*(part.meshSize().y()+1),
					2*part.meshSize().prod() + part.meshSize().y() - 1
				) ).all() );
		else
			REQUIRE( part.mesh().boundaryPatch("left").size() == 0 );
		/* ad libitum */
	};
	part.comm( comm );

	MESH_TYPE mesh;

	if ( comm.isMaster() ){
		mesh = uniMeshAs<MESH_TYPE>(
			globalMeshSize, (Array2s{1,1}), (Array2s{0,0}), nGhostLayers );
		addStandardBoundaryPatches(mesh);
	}
	part.assignAndDistribute( std::move(mesh) );
	// std::cout << "nRanks = " << comm.size()
	// 	<< ", as " << part.nParts().x() << " x " << part.nParts().y() << "\n";
	// print( part.meshSize().transpose() );
	
	test();

	/* let's pretend the meshes had been written to file, read back in and
	 * were locally available for each rank */
	mesh = uniMeshAs<MESH_TYPE>( part.meshSize(), (Array2s{1,1}), (Array2s{0,0}) );
	mesh.addBoundaryPatches( part.mesh().moveBoundaryPatches() );
	part = { std::move(mesh), part.offset(), comm };

	test();

	/* redistribution */
	Array2i newSize { part.meshSize() - 1 };
	if ( part.indexInBand() == part.nParts().x() - 1 )
		newSize.x() += part.nParts().x();
	if ( part.band() == part.nParts().y() - 1 )
		newSize.y() += part.nParts().y();

	Array2Xi redistSizes ( 2, comm.size() );
	MPI_Allgather(
		newSize.data(), 2, MPI_INDEX,
		redistSizes.data(), 2, MPI_INDEX, comm
	);

	part.redistributeMesh(redistSizes);
}

#undef MESH_NAME
#undef MESH_TYPE